import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import matplotlib.patches as patches
import os
from matplotlib.colors import LinearSegmentedColormap
from datetime import datetime

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

# 定义常量 - 保存结果的目录
RESULTS_DIR = r"D:\大三上\大数据分析及数据可视化\《Excel数据可视化 - 从图表到数据大屏》-清华-郭宏远\实验\results"


def load_erp_data():
    """加载ERP订单数据"""
    # 定义可能的文件路径
    possible_paths = [
        r"D:\大三上\大数据分析及数据可视化\《Excel数据可视化 - 从图表到数据大屏》-清华-郭宏远\实验\data\erp_order_data.xlsx",
        r"D:\大三上\大数据分析及数据可视化\《Excel数据可视化 - 从图表到数据大屏》-清华-郭宏远\实验\erp_order_data.xlsx"
    ]

    for path in possible_paths:
        if os.path.exists(path):
            try:
                df = pd.read_excel(path)
                # 确保日期格式正确
                if 'order_time' in df.columns:
                    df['order_time'] = pd.to_datetime(df['order_time'])
                print(f"✓ 数据加载成功: {path}")
                return df
            except Exception as e:
                print(f"读取 {path} 失败: {e}")
                continue

    raise FileNotFoundError("未找到erp_order_data.xlsx文件，请确认文件是否存在或路径是否正确")


def process_region_data(df):
    """处理区域数据，对比最近两个完整年份的销量"""
    # 确保order_time是datetime类型
    if 'order_time' not in df.columns:
        raise ValueError("数据中缺少'order_time'列，无法进行时间分析")

    # 确保province列存在
    if 'province' not in df.columns:
        raise ValueError("数据中缺少'province'列，无法进行区域分析")

    # 确保product_amount列存在
    if 'product_amount' not in df.columns:
        raise ValueError("数据中缺少'product_amount'列，无法计算销售金额")

    # 检查数据中的年份范围
    years = df['order_time'].dt.year.unique()
    years = sorted(years)
    print(f"数据中包含的年份: {years}")

    # 选择最近的两个完整年份
    if len(years) < 2:
        raise ValueError("数据中年份不足，需要至少两年的数据进行对比")

    # 选择最新的两个年份
    year1 = years[-2]  # 倒数第二年
    year2 = years[-1]  # 最新年份

    # 为今年选择最近6个月
    latest_date = df['order_time'].max()
    current_year = latest_date.year
    current_month = latest_date.month

    # 如果是完整年份，使用全年数据，否则使用最近6个月
    if current_month >= 12:  # 假设12月表示全年数据
        data_year2 = df[df['order_time'].dt.year == year2]
        period_desc = '全年'
    else:
        # 使用最近6个月
        start_date = latest_date - pd.DateOffset(months=5)
        data_year2 = df[(df['order_time'] >= start_date) & (df['order_time'] <= latest_date)]
        period_desc = f'{start_date.month}月-{latest_date.month}月'

    # 选择去年的同期数据
    last_year_start = pd.Timestamp(year=year1, month=start_date.month,
                                   day=1) if 'start_date' in locals() else pd.Timestamp(year=year1, month=1, day=1)
    last_year_end = pd.Timestamp(year=year1, month=latest_date.month, day=latest_date.day)
    data_year1 = df[(df['order_time'].dt.year == year1) &
                    (df['order_time'].dt.month >= last_year_start.month) &
                    (df['order_time'].dt.month <= last_year_end.month)]

    print(f"使用 {year1}年 对比 {year2}年{period_desc}的数据")
    print(f"{year1}年记录数: {len(data_year1)}")
    print(f"{year2}年{period_desc}记录数: {len(data_year2)}")

    # 按地区汇总销量
    regions_year1 = data_year1.groupby('province')['product_amount'].sum().reset_index()
    regions_year2 = data_year2.groupby('province')['product_amount'].sum().reset_index()

    # 重命名列
    regions_year1.columns = ['province', f'sales_{year1}']
    regions_year2.columns = ['province', f'sales_{year2}']

    # 合并两个数据集
    merged = pd.merge(regions_year1, regions_year2, on='province', how='inner')

    # 只保留主要地区（华南、华北、东北、西北、华东）
    main_regions = ['华南', '华北', '东北', '西北', '华东', '北京市', '上海市', '广东省', '江苏省', '浙江省', '四川省',
                    '湖北省', '陕西省']
    available_regions = [r for r in main_regions if r in merged['province'].values]

    if not available_regions:
        # 如果主要地区都没有，使用数据中存在的所有地区
        print("警告：数据中没有包含主要地区。将使用所有可用地区。")
        available_regions = merged['province'].tolist()

    merged = merged[merged['province'].isin(available_regions)]

    # 按最近年份销售量排序（从高到低）
    merged = merged.sort_values(f'sales_{year2}', ascending=False)

    # 重命名列以便后续处理
    merged = merged.rename(columns={
        f'sales_{year1}': 'sales_last',
        f'sales_{year2}': 'sales_current'
    })

    # 存储年份信息
    merged.attrs['year1'] = year1
    merged.attrs['year2'] = year2
    merged.attrs['period_desc'] = period_desc

    return merged


def create_butterfly_chart():
    """创建区域销量对比蝴蝶图（左右对称分布）"""
    # 加载数据
    df = load_erp_data()

    # 处理区域数据
    region_data = process_region_data(df)

    # 提取数据
    regions = region_data['province'].tolist()
    sales_last = region_data['sales_last'].tolist()
    sales_current = region_data['sales_current'].tolist()

    # 获取年份信息
    year1 = region_data.attrs['year1']
    year2 = region_data.attrs['year2']
    period_desc = region_data.attrs['period_desc']

    # 检查数据有效性
    if not regions:
        print("错误：没有找到有效的区域数据。")
        return None, None

    # 创建渐变颜色映射
    blue_colors = ['#16213E', '#0F3460', '#1A5F7A', '#2C8C99', '#4BB5C2', '#7FD6E3']
    red_colors = ['#5C1E2B', '#8A2D3D', '#B73E4F', '#D54D5F', '#E56A72', '#F08B8F']

    # 创建颜色映射
    blue_cmap = LinearSegmentedColormap.from_list('blue', blue_colors, N=100)
    red_cmap = LinearSegmentedColormap.from_list('red', red_colors, N=100)

    # 创建图形 - 靛蓝色背景
    fig = plt.figure(figsize=(14, 9), facecolor='#1A1A2E')
    ax = fig.add_subplot(111, facecolor='#1A1A2E')

    # 设置y轴位置
    y_pos = np.arange(len(regions))

    # 设置条形高度
    bar_height = 0.4

    # 创建去年数据条 - 左侧
    for i, (y, sale) in enumerate(zip(y_pos, sales_last)):
        gradient_width = 80
        for j in range(gradient_width):
            width_segment = sale / gradient_width
            x_left = -j * width_segment
            color = red_cmap(1 - (j / gradient_width) * 0.7)
            rect = patches.Rectangle(
                (x_left, y - bar_height / 2), width_segment, bar_height,
                linewidth=0, facecolor=color, alpha=0.9
            )
            ax.add_patch(rect)

    # 创建今年数据条 - 右侧
    for i, (y, sale) in enumerate(zip(y_pos, sales_current)):
        gradient_width = 80
        for j in range(gradient_width):
            width_segment = sale / gradient_width
            x_left = j * width_segment
            color = blue_cmap(1 - (j / gradient_width) * 0.7)
            rect = patches.Rectangle(
                (x_left, y - bar_height / 2), width_segment, bar_height,
                linewidth=0, facecolor=color, alpha=0.9
            )
            ax.add_patch(rect)

    # 添加白色边框
    for i, y in enumerate(y_pos):
        # 去年边框
        if sales_last[i] > 0:
            ax.plot([-sales_last[i], 0], [y - bar_height / 2, y - bar_height / 2], color='#E0E0E0', linewidth=1.5)
            ax.plot([-sales_last[i], -sales_last[i]], [y - bar_height / 2, y + bar_height / 2], color='#E0E0E0',
                    linewidth=1.5)
            ax.plot([-sales_last[i], 0], [y + bar_height / 2, y + bar_height / 2], color='#E0E0E0', linewidth=1.5)

        # 今年边框
        if sales_current[i] > 0:
            ax.plot([0, sales_current[i]], [y - bar_height / 2, y - bar_height / 2], color='#E0E0E0', linewidth=1.5)
            ax.plot([sales_current[i], sales_current[i]], [y - bar_height / 2, y + bar_height / 2], color='#E0E0E0',
                    linewidth=1.5)
            ax.plot([0, sales_current[i]], [y + bar_height / 2, y + bar_height / 2], color='#E0E0E0', linewidth=1.5)

    # 添加数据标签
    for i, (sale_last, sale_current) in enumerate(zip(sales_last, sales_current)):
        # 去年标签 - 左侧
        if sale_last > 0:
            ax.text(-sale_last * 0.95, i, f'{int(sale_last)}',
                    ha='right', va='center', fontsize=12, fontweight='bold', color='#FFFFFF',
                    bbox=dict(boxstyle='round,pad=0.2', facecolor='#5C1E2B', edgecolor='none', alpha=0.7))

        # 今年标签 - 右侧
        if sale_current > 0:
            ax.text(sale_current * 0.95, i, f'{int(sale_current)}',
                    ha='left', va='center', fontsize=12, fontweight='bold', color='#FFFFFF',
                    bbox=dict(boxstyle='round,pad=0.2', facecolor='#2C3E50', edgecolor='none', alpha=0.7))

    # 添加年份标签
    max_sales = max(max(sales_last), max(sales_current))
    ax.text(0, -1.2, f'{year1}',
            ha='center', va='center', fontsize=14, fontweight='bold', color='#E56A72')
    ax.text(0, len(regions), f'{year2}年{period_desc}',
            ha='center', va='center', fontsize=14, fontweight='bold', color='#4BB5C2')

    # 设置坐标轴
    ax.set_yticks(y_pos)
    ax.set_yticklabels(regions, fontsize=14, fontweight='bold', color='#E0E0E0')
    ax.tick_params(axis='x', labelsize=14, colors='#E0E0E0')
    ax.set_xticks([])

    # 添加垂直中心线
    ax.axvline(x=0, color='#E0E0E0', linestyle='-', linewidth=1.5)

    # 设置标题
    title_main = f'{year2}年{period_desc}各区域对比{year1}年同期销量'

    # 分析最近的趋势
    total_last = sum(sales_last)
    total_current = sum(sales_current)
    growth_rate = ((total_current - total_last) / total_last * 100) if total_last > 0 else 0

    declining_regions = []
    for i, region in enumerate(regions):
        if sales_last[i] > 0 and sales_current[i] < sales_last[i]:
            declining_regions.append(region)

    if declining_regions:
        subtitle = f'{year2}年{period_desc}整体销量{"增长" if growth_rate > 0 else "下降"}{abs(growth_rate):.1f}%，{"、".join(declining_regions)}区域较{year1}年同期有所下降'
    else:
        subtitle = f'{year2}年{period_desc}整体销量{"增长" if growth_rate > 0 else "下降"}{abs(growth_rate):.1f}%，所有区域均较{year1}年同期有所{"增长" if growth_rate > 0 else "下降"}'

    # 修复标题重叠问题 - 关键修改：调整y参数，使标题向上移动
    ax.set_title(f"{title_main}\n{subtitle}",
                 fontsize=20, fontweight='bold',
                 # 调整y参数，将标题向上移动
                 y=1.08,  # 关键修改：将y值从默认1.0增加到1.08
                 pad=25,
                 color='#FFFFFF')

    # 美化坐标轴
    for spine in ax.spines.values():
        spine.set_color('#4A4A6A')
        spine.set_linewidth(2)

    # 设置网格线
    ax.grid(axis='x', alpha=0.3, linestyle='--', color='#4A4A6A')
    ax.set_axisbelow(True)

    # 设置x轴范围
    max_sales = max(max(sales_last), max(sales_current)) * 1.15
    ax.set_xlim(-max_sales, max_sales)

    # 添加数据来源
    latest_date = df['order_time'].max()
    latest_date_str = latest_date.strftime('%Y.%m.%d')
    source_text = f'*注：数据来源于公司销售系统，统计日期截至{latest_date_str}'
    ax.text(0.02, 0.02, source_text, transform=ax.transAxes,
            fontsize=10, color='#B0B0B0', alpha=0.7, va='bottom')

    # 调整布局
    plt.tight_layout()

    # 确保结果目录存在
    os.makedirs(RESULTS_DIR, exist_ok=True)
    # 保存图片
    output_path = os.path.join(RESULTS_DIR, '06_蝴蝶图.png')
    plt.savefig(output_path, dpi=300, bbox_inches='tight',
                facecolor='#1A1A2E', edgecolor='none')

    plt.show()

    # 数据分析
    print("\n区域销量对比图数据分析：")
    print(f"- 数据覆盖区域：{len(regions)}个地区")
    print(f"- {year2}年{period_desc}总销量：{total_current:.2f}万元")
    print(f"- {year1}年同期总销量：{total_last:.2f}万元")
    print(f"- 整体同比增长率：{growth_rate:.1f}%")

    if declining_regions:
        print(f"- 销量下降的区域：{', '.join(declining_regions)}")

    # 找出增长最快的区域
    growth_rates = []
    for i, region in enumerate(regions):
        if sales_last[i] > 0:
            growth = ((sales_current[i] - sales_last[i]) / sales_last[i]) * 100
            growth_rates.append((region, growth))

    if growth_rates:
        fastest_growing = max(growth_rates, key=lambda x: x[1])
        slowest_growing = min(growth_rates, key=lambda x: x[1])
        print(f"- 增长最快的区域：{fastest_growing[0]}，增长率{fastest_growing[1]:.1f}%")
        print(f"- 增长最慢的区域：{slowest_growing[0]}，增长率{slowest_growing[1]:.1f}%")

    return fig, ax


# 执行代码
if __name__ == "__main__":
    try:
        fig, ax = create_butterfly_chart()
    except Exception as e:
        print(f"图表生成失败: {e}")
        raise
